import argparse
import os
import torch
import cv2
import numpy as np
from PIL import Image
from diffusers import StableDiffusionPipeline
from src.steganography import Steganography
import pandas as pd  # 导入 pandas 库

device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')

# 从本地加载模型 stable diffusion v1.5
local_model_path = '/root/autodl-tmp/stable-diffusion-v1-5'
ldm_stable = StableDiffusionPipeline.from_pretrained(local_model_path).to(device)

def load_image(image_path):
    image = Image.open(image_path).convert("RGB")
    h, w = image.size
    image = np.array(image.resize((512, 512)))  # 注意这里要 resize 图像大小
    return image, h, w

if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument('--csv_path', type=str, default='prompt.csv', help='path to the CSV file containing image and text prompts')
    parser.add_argument('--save_path', type=str, default='./output', help='output save path')
    args = parser.parse_args()

    if not os.path.exists(args.save_path):
        os.makedirs(args.save_path)

    # 读取 CSV 文件
    df = pd.read_csv(args.csv_path)

    # 检查所需的列是否存在
    if not all(col in df.columns for col in ['Photo', 'private_key', 'public_key']):
        raise ValueError("CSV文件必须包含 'Photo', 'private_key', 'public_key' 三列")

    steg = Steganography(ldm_stable, device)

    for index, row in df.iterrows():
        image_path = row['Photo']
        private_key = row['private_key']
        public_key = row['public_key']

        image_name_without_ext = os.path.splitext(os.path.basename(image_path))[0]  

        # 组合图像路径以便加载
        full_image_path = os.path.join('./data', image_path)

        image_gt, h, w = load_image(full_image_path)
        
        current_output_dir = os.path.join(args.save_path, image_name_without_ext)
        os.makedirs(current_output_dir, exist_ok=True)

        # 保存原始图像
        cv2.imwrite("{:s}/gt.png".format(current_output_dir), cv2.cvtColor(image_gt, cv2.COLOR_RGB2BGR))

        # 隐藏过程
        latent_noise = steg.invert(private_key, steg.image2latent(image_gt), is_forward=True)  # 正向操作
        image_hide_latent = steg.invert(public_key, latent_noise, is_forward=False)  # 反向操作

        # 保存隐藏图片
        image_hide = steg.latent2image(image_hide_latent)
        cv2.imwrite("{:s}/hide.png".format(current_output_dir), cv2.cvtColor(image_hide, cv2.COLOR_RGB2BGR))
        
        # 还原过程
        image_hide_latent_reveal = steg.image2latent(image_hide)

        latent_noise = steg.invert(public_key, image_hide_latent_reveal, is_forward=True)
        image_reverse_latent = steg.invert(private_key, latent_noise, is_forward=False)

        image_reverse = steg.latent2image(image_reverse_latent)
        cv2.imwrite("{:s}/reverse.png".format(current_output_dir), cv2.cvtColor(image_reverse, cv2.COLOR_RGB2BGR))
        
        print(index + 1 , ".jpg END =========================================================================")
